from math import log
import operator


def createDataSet():
    # dataSet = [[1, 1, 'yes'],
    #            [1, 1, 'yes'],
    #            [1, 0, 'no'],
    #            [0, 1, 'no'],
    #            [0, 1, 'no']]
    # labels = ['no surfacing', 'flippers']
    dataSet = [[1, 1, 1, 'yes'],
               [1, 1, 1, 'yes'],
               [1, 1, 0, 'no'],
               [1, 0, 1, 'no'],
               [0, 1, 1, 'no'],
               [0, 1, 1, 'no']]
    labels = ['no surfacing', 'flippers', 'other']
    # change to discrete values
    return dataSet, labels


# 2.1
def calcShannonEnt(dataSet):
    dataLen = len(dataSet)
    labelCount = {}
    for item in dataSet:
        curLabel = item[-1]
        if curLabel not in labelCount.keys():
            labelCount[curLabel] = 0
        labelCount[curLabel] += 1
    shannonEnt = 0.0
    for item in labelCount:
        prob = float(labelCount[item]) / dataLen
        shannonEnt -= prob * log(prob, 2)
    return shannonEnt


# 2.2
# 1:select feature i, 2:select feature value
def spliteDataSet(dataSet, feat, value):
    subData = []
    for item in dataSet:
        if item[feat] == value:
            subItem = item[:feat]
            subItem.extend(item[feat + 1:])
            subData.append(subItem)
    return subData


# splite dataset,   calulate info gain,   choose best feature by info gain max
# return column index
def choseBestFeatureToSplit(dataSet):
    features = len(dataSet[0]) - 1
    baseEnt = calcShannonEnt(dataSet)  # calulate base Ent
    initInfoGain = 0
    bestFeature = -1
    for featurei in range(features):
        featList = [item[featurei] for item in dataSet]  # get axis line data
        featListUniq = set(featList)
        subEnt = 0.0
        for featVal in featListUniq:
            subDataset = spliteDataSet(dataSet, featurei, featVal)
            prop = len(subDataset) / float(len(dataSet))
            subEnt += prop * calcShannonEnt(subDataset)
        infoGain = baseEnt - subEnt  # calulate infogain
        if initInfoGain < infoGain:
            initInfoGain = infoGain
            bestFeature = featurei  # choose infoGain best feature
    return bestFeature


def majorityCnt(classList):
    classCount = {}
    for vote in classList:
        if vote not in classCount.keys(): classCount[vote] = 0
        classCount[vote] += 1
    sortedClassCount = sorted(classCount.iteritems(), key=operator.itemgetter(1), reverse=True)
    return sortedClassCount[0][0]


def createTree(dataSet, label):
    classList = [item[-1] for item in dataSet]  # get result column
    if classList.count(classList[0]) == len(classList):
        return classList[0]  # return result
    if len(dataSet[0]) == 1:  # stop splitting when there are no more features in dataSet
        return majorityCnt(classList)
    bestFeatureIndex = choseBestFeatureToSplit(dataSet)
    bestFeature = label[bestFeatureIndex]
    tree = {bestFeature: {}}
    del (label[bestFeatureIndex])
    featVals = [item2[bestFeatureIndex] for item2 in dataSet]
    uniFeatVals = set(featVals)
    for featVal in uniFeatVals:
        subLabel = label[:]
        subData = spliteDataSet(dataSet, bestFeatureIndex, featVal)
        subTree = createTree(subData, subLabel)  # recursion get subTree or result
        tree[bestFeature][featVal] = subTree
    return tree


def classify(tree, labels, data):
    firstFeat = tree.keys()[0]
    subDict = tree[firstFeat]
    firstFeatIndex = labels.index(firstFeat)
    key = data[firstFeatIndex]
    sub = subDict[key]
    if isinstance(sub, dict):
        return classify(sub, labels, data)
    else:
        return sub


def storeTree(tree, file='tree.txt'):
    import pickle
    fw = open(file, 'w')
    pickle.dump(tree, fw)
    fw.close()


def getTreeByFile(file='tree.txt'):
    import pickle
    fw = open(file)
    return pickle.load(fw)


def getLensesData():
    fw = open('lenses.txt')
    lenses = [item.strip().split('\t') for item in fw.readlines()]
    # feature
    labels = ['age', 'prescript', 'astigmatic', 'tearRate']
    # feature value
    # hard,soft, no lenses
    return lenses, labels


# dataSet, label = createDataSet()
dataSet, label = getLensesData()
label1 = label[:]
# tree=getTreeByFile()
tree = createTree(dataSet, label)
# print tree
# storeTree(tree)
print classify(tree, label1, [0, 1])
# print classify(tree, label1, ['pre', 'myope', 'no', 'normal'])
# print tree.keys()
# print 'over'
# bestFeature = choseBestFeatureToSplit(dataSet)
# print bestFeature
# shannonEnt = calcShannonEnt(dataSet)
# print shannonEnt
